#include <iomanip>
#include <cmath>
#include <cstring>
#include <lapacke.h>
#include "FiniteDiff/LevelOp.h"
#include "FiniteDiff/GhostFiller.h"
#include "Embedded/PoissonOp.h"
#include "Embedded/EmFuncFiller.h"
#include "Embedded/SourceStencil.impl.h"
#include "Embedded/EvaluateMonomials.h"
#include "Embedded/ConditionNumber.h"
#include "../example/common/TestUtility.h"

template <int Dim>
struct PoissonOp<Dim>::LaplacianDescriptor
{
  enum { fitOrder = PoissonOp<Dim>::fitOrder };

  bool checkStandardScheme(const Tensor<int, Dim> &nodeLabels, Vec<int, Dim> idx, bool useOuterRect) const
  {
    auto validBox = nodeLabels.box();
    for(int d = 0; d < Dim; ++d) {
      if (idx[d] < validBox.lo()[d] + 2 || idx[d] > validBox.hi()[d] - 2)
        return useOuterRect;
    }
    for(int d = 0; d < Dim; ++d) {
      auto ej = iVec::unit(d);
      for(int k = -2; k <= 2; ++k)
        if(nodeLabels(idx + ej*k) < 0)
          return false;
    }
    return true;
  }

  template <class... Ts>
  Real transform(const Vec<Real, Dim> &p, Ts... indices) const {
    Vec<int, Dim> beta { decltype(indices)::value ... };
    rVec partial2nd;
    for(int d = 0; d < Dim; ++d) {
      if(beta[d] < 2) {
        partial2nd[d] = 0.0;
        continue;
      }
      partial2nd[d] = 1.0;
      beta[d] -= 2;
      for(int k = 0; k < Dim; ++k)
        partial2nd[d] *= std::pow(p[k], beta[k]);
      beta[d] += 2;
      partial2nd[d] *= beta[d] * (beta[d]-1);
    }
    return sum(partial2nd);
  }
};

//======================================================================

template <int Dim>
PoissonOp<Dim>::PoissonOp(const PoissonOp::EmGrid &aGrid, const char *aBcTypes)
    : emGrid(aGrid)
{
  strncpy(bcTypes, aBcTypes, 2*Dim+2);
  // 1. Construct the stencils.
  LaplacianDescriptor descriptor;
  SourceStencilParam<Dim> ssparam;
  ssparam.otype = SourceStencilParam<Dim>::OperatorType::Elliptic;
  CPUTimer timer;
  timer.reset();
  SourceStencil<Dim> ss(emGrid, descriptor, ssparam);
  std::cout << "Construction of stencils : " << std::defaultfloat << std::setprecision(5)
            << timer() << "s. " << std::endl;
#ifndef NDEBUG
  if(get_dbglevel() >= 2)
    ss.dump(tmpos);
#endif
  // 2. Find the approximation.
  const auto &bdryPoints = emGrid.bdryPoints;
  const auto &outNormals = emGrid.outNormals;
  const auto dx = emGrid.spacing();
  const int n = fitOrder - 1;
  const int nBasis = binom(n+Dim, Dim);
  auto DirichletModifier = [&](const Real &x, const auto &) { return x; };
  auto NeumannModifier = [&](const rVec &x, const rVec &n) { return dot(x, n); };

  timer.reset();
  for(const auto &stencil : ss.vStencil) {
    const int nSources = stencil.locations.size();
    const auto destIdx = stencil.destIdx;
    const auto refLoc  = (destIdx + 0.5) * dx;
    // 2.1. calculate the sample matrix.
    Tensor<Real, 2> sampleMatrix(Vec<int, 2> {nBasis, nSources});
    sampleMatrix = 0.0;
    Real *where2Fill = sampleMatrix.data();
    for(int k = 0; k < nSources; ++k) {
      rVec p, nl;
      if(stencil.isBdry[k] == 0) {
        iVec jdx = stencil.locations[k];
        p = (jdx + 0.5) * dx;
        EvaluateMonomials<Dim,n>::doEvaluate(identityOfMonomials(), DirichletModifier, where2Fill, p-refLoc, rVec(0.0), dx);
      } else {
        // use the nearest boudary point as the boundary condition
        p = bdryPoints(destIdx);
        nl = outNormals(destIdx);
        if(aBcTypes[2*Dim] == 'N') {
//          stencil.weights[k] *= sqrt(norm(dx, 2));
          EvaluateMonomials<Dim,n>::doEvaluate(gradOfMonomials(), NeumannModifier, where2Fill, p-refLoc, nl, dx);
        } else {
          EvaluateMonomials<Dim,n>::doEvaluate(identityOfMonomials(), DirichletModifier, where2Fill, p-refLoc, rVec(0.0), dx);
        }
      }
    }
    assert((where2Fill - sampleMatrix.data()) == sampleMatrix.volume());
    // scale the columns corresponding to the Neumann condition
    const Real regularizer4Neumann = min_of(dx);
    for(int k = 0; k < nSources; ++k) {
      if(stencil.isBdry[k] && aBcTypes[2*Dim] == 'N') {
        for (int j = 0; j < nBasis; ++j)
          sampleMatrix(j, k) *= regularizer4Neumann;
      }
    }
#ifndef NDEBUG
    const Real cn = estimateUnderdetCond(sampleMatrix);
    if(cn >= dangerousCN || get_dbglevel() >= 2)
      dbgcout1 << "The linear system for " << destIdx
               << " has condition number " << std::setprecision(2) << std::scientific << cn
               << ". \n" << std::defaultfloat;
    dbgcout3 << std::showpos << std::setprecision(3) << sampleMatrix << "\n";
    dbgcout3 << std::defaultfloat;
#endif
    // 2.2. calculate the RHS.
    Tensor<Real, 1> sampleRHS(nSources);
    sampleRHS = 0.0;
    where2Fill = sampleRHS.data();
    EvaluateMonomials<Dim,n>::doEvaluate(descriptor, DirichletModifier, where2Fill, rVec(0.0), rVec(0.0), dx);
    // 2.3. solve the linear system to obtain the coefficients.
    dbgcout3 << sampleMatrix << std::endl;
    // scale the columns
    for(int j = 0; j < nSources; ++j)
      for(int i = 0; i < nBasis; ++i)
        sampleMatrix(i, j) /= stencil.weights[j];
    auto info = LAPACKE_dgels(LAPACK_COL_MAJOR, 'N'/*trans*/, nBasis/*m*/, nSources/*n*/, 1/*nrhs*/,
                              sampleMatrix.data()/*a*/, nBasis/*lda*/, sampleRHS.data()/*b*/, nSources/*ldb*/);
    if(info != 0)
      throw std::runtime_error("[PoissonOp::PoissonOp] DGELS failed.");
    // scale the raw solution
    for(int j = 0; j < nSources; ++j) {
      sampleRHS(j) /= stencil.weights[j];
      if(stencil.isBdry[j] && aBcTypes[2*Dim] == 'N')
        sampleRHS(j) *= regularizer4Neumann;
    }
    // 2.4. Save the coefficients in sparse matrix format.
    matOp.insertRow(destIdx, nSources, stencil.locations.data(), sampleRHS.data());
    dbgcout2 << sampleRHS << "\n";
  }
  std::cout << "Discretization of the operator : " << std::defaultfloat << std::setprecision(5)
            << timer() << "s. " << std::endl;
}

template <int Dim>
void PoissonOp<Dim>::apply(const Tensor<Real, Dim> &aData, Tensor<Real, Dim> &outData) const
{
  // 1. Apply the standard scheme.
  LevelOp<Dim> lvOp(emGrid);
  lvOp.LAPC_4TH(aData, outData);
  // 2. Apply the correction.
  matOp.AXPY(1.0, aData, [](iVec) { return 0.0; }, outData);
}

template <int Dim>
void PoissonOp<Dim>::apply(const Tensor<Real, Dim> &aData, const Tensor<Real, Dim> &rhs, Tensor<Real, Dim> &rsd) const
{
  // 1. Apply the standard scheme.
  LevelOp<Dim> lvOp(emGrid);
  lvOp.LAPC_4TH(aData, rhs, rsd);
  // 2. Apply the correction.
  matOp.AXPY(-1.0, aData, rhs, rsd);
}

template class PoissonOp<SpaceDim>;

//======================================================================
template <int Dim>
PoissonBottomSolver<Dim>::PoissonBottomSolver(const PoissonOp<Dim> &aPoissonOp)
    : emGrid(aPoissonOp.emGrid), poissonOp(aPoissonOp)
{
  CPUTimer timer;
  timer.reset();
  bcTypes = poissonOp.bcTypes;
  Box<Dim> gbx = emGrid.getGhostedBox();
  unknownMap.resize(gbx);
  unknownMap = -1;
  numUnknowns = 0;
  // ===================================================================
  // find the one-to-one correspondences between the cells and the unknowns
  const auto &nodeLabels = emGrid.nodeLabels;
//  using NodeType = EmbeddedGrid<Dim>::NodeType;
  ddfor(emGrid, [&](const iVec &idx) {
    if(nodeLabels(idx) >= 0)
      unknownMap(idx) = numUnknowns++;
  }, false);
  // when useOuterRect, handle the ghost cells
  const int nG = emGrid.getNumGhost();
  if(emGrid.useOuterRect) {
    for (int d = 0; d < Dim; ++d) {
      Box<Dim-1> sideBox = reduce(emGrid, d);
      for (int side : {-1, 1}) {
        int a = (side == -1) ? (emGrid.lo()[d]) : (emGrid.hi()[d]);
        a += side;
        for (int k = 0; k < nG; ++k, a += side) {
          ddfor(sideBox, [&](const Vec<int, Dim-1> &idx) {
            unknownMap(enlarge(idx, a, d)) = numUnknowns++;
          }, false);
        } // end loop of each layer of ghosts
      } // end loop of lo/hi side
    } // end loop of dimensions
  }
  // ===================================================================
  // fill in the coefficients of the interior cells and the irregular cells
  auto dx = emGrid.spacing();
  poissonMat.resize(numUnknowns+1);
  poissonMat = 0.0;
  const auto &sp = poissonOp.matOp;
  auto rit = sp.cbegin();
  VecCompare<int, Dim> ivcmp;
  ddfor(emGrid, [&](const iVec &idx) {
    if(unknownMap(idx) == -1)
      return;
    // check if the coefficients for this row have been generated
    if(rit != sp.cend()) {
      int cmp = ivcmp.compare(idx, rit.getRow());
      assert(cmp <= 0);
      if(cmp == 0) {
        int m0 = unknownMap(idx);
        for(int k = 0; k < rit.getNz(); ++k) {
          auto col = rit.getColumn(k);
          int m1 = unknownMap(col);
//            assert(m1 != -1);
          if(m1 >= 0)
            poissonMat(m0, m1) = rit.getValue(k);
        }
        ++rit;
        return;
      }
    }
    // use the standard difference scheme
    const int i = idx[0], j = idx[1], k = (Dim >= 3) ? (idx[2]) : (0);
    int m[13];
    if(Dim == 2) {
      m[0] = unknownMap(i, j);
      m[1] = unknownMap(i-2, j); m[2] = unknownMap(i-1, j);
      m[3] = unknownMap(i+1, j); m[4] = unknownMap(i+2, j);
      m[5] = unknownMap(i, j-2); m[6] = unknownMap(i, j-1);
      m[7] = unknownMap(i, j+1); m[8] = unknownMap(i, j+2);
    } else if(Dim == 3) {
      m[0] = unknownMap(i, j, k);
      m[1] = unknownMap(i-2, j, k); m[2] = unknownMap(i-1, j, k);
      m[3] = unknownMap(i+1, j, k); m[4] = unknownMap(i+2, j, k);
      m[5] = unknownMap(i, j-2, k); m[6] = unknownMap(i, j-1, k);
      m[7] = unknownMap(i, j+1, k); m[8] = unknownMap(i, j+2, k);
      m[9] = unknownMap(i, j, k-2); m[10] = unknownMap(i, j, k-1);
      m[11] = unknownMap(i, j, k+1); m[12] = unknownMap(i, j, k+2);
    }
    poissonMat(m[0], m[0]) = -30.0/12 * (1.0/(dx[0]*dx[0]) + 1.0/(dx[1]*dx[1]));
    if(Dim >= 3)
      poissonMat(m[0], m[0]) += -30.0/12 / (dx[2]*dx[2]);
    poissonMat(m[0], m[1]) = poissonMat(m[0], m[4]) = -1.0/12 / (dx[0]*dx[0]);
    poissonMat(m[0], m[2]) = poissonMat(m[0], m[3]) = 16.0/12 / (dx[0]*dx[0]);
    poissonMat(m[0], m[5]) = poissonMat(m[0], m[8]) = -1.0/12 / (dx[1]*dx[1]);
    poissonMat(m[0], m[6]) = poissonMat(m[0], m[7]) = 16.0/12 / (dx[1]*dx[1]);
    if(Dim >= 3) {
      poissonMat(m[0], m[9]) = poissonMat(m[0], m[12]) = -1.0 / 12 / (dx[2] * dx[2]);
      poissonMat(m[0], m[10]) = poissonMat(m[0], m[11]) = 16.0 / 12 / (dx[2] * dx[2]);
    }
  }, false);
  // when useOuterRect, fill in the coefficients for the ghost cells (hardcoded to two layers of ghosts)
  if(emGrid.useOuterRect) {
    for (int d = 0; d < Dim; ++d) {
      Box<Dim-1> sideBox = reduce(emGrid, d);
      for (int side : {-1, 1}) {
        int a = (side == -1) ? (emGrid.lo()[d]) : (emGrid.hi()[d]);
        ddfor(sideBox, [&](const Vec<int, Dim-1> &idx) {
          iVec kdx = enlarge(idx, a, d);
          auto u = iVec::unit(d) * side;
          int m[] = {unknownMap(kdx + u * 2), unknownMap(kdx + u), unknownMap(kdx),
                     unknownMap(kdx - u), unknownMap(kdx - u * 2), unknownMap(kdx - u * 3)};
          switch(bcTypes[2*d + (side+1)/2]) {
          case 'N':
            // interpolate using the Neumann data to the near ghost
            poissonMat(m[1], m[1]) = -1.0;
            poissonMat(m[1], m[2]) = 17.0 / 22;
            poissonMat(m[1], m[3]) = 9.0 / 22;
            poissonMat(m[1], m[4]) = -5.0 / 22;
            poissonMat(m[1], m[5]) = 1.0 / 22;
            // interpolate using the Neumann data to the far ghost
            poissonMat(m[0], m[0]) = -1.0;
            poissonMat(m[0], m[2]) = -135.0 / 22;
            poissonMat(m[0], m[3]) = 265.0 / 22;
            poissonMat(m[0], m[4]) = -135.0 / 22;
            poissonMat(m[0], m[5]) = 27.0 / 22;
            break;
          case 'D':
            //
            poissonMat(m[1], m[1]) = -1.0;
            poissonMat(m[1], m[2]) = -4.0;
            poissonMat(m[1], m[3]) = 2.0;
            poissonMat(m[1], m[4]) = -4.0 / 5;
            poissonMat(m[1], m[5]) = 1.0 / 7;
            //
            poissonMat(m[0], m[0]) = -1.0;
            poissonMat(m[0], m[2]) = -30.0;
            poissonMat(m[0], m[3]) = 20.0;
            poissonMat(m[0], m[4]) = -9.0;
            poissonMat(m[0], m[5]) = 12.0 / 7;
            break;
          }
        }, false);
      } // end loop of lo/hi side
    } // end loop of dimesnions
  }
  // ===================================================================
  // modify the matrix to handle the singularity of the Neumann problem
  pureNeumann = (strchr(bcTypes, 'D') == nullptr);
  numActualUnknowns = numUnknowns + pureNeumann;
  if(pureNeumann) {
    for (int k = 0; k < numUnknowns; ++k)
      poissonMat(k, numUnknowns) = -1.0;
    for (int k = 0; k < numUnknowns; ++k)
      poissonMat(numUnknowns, k) = 1.0;
  }
  // ===================================================================
  // cache the factorization of the Poisson matrix
  poissonFact = poissonMat.copy();
  ipiv.resize(numUnknowns + 1);
  auto info = LAPACKE_dgetrf(LAPACK_COL_MAJOR, numActualUnknowns/*m*/, numActualUnknowns/*n*/,
                             poissonFact.data()/*a*/, numUnknowns + 1/*lda*/, ipiv.data()/*ipiv*/);
  if(info != 0)
    throw std::runtime_error("[DirectSolver::DirectSolver] DGETRF fails.");
  std::cout << "Factorization of the bottom solver : " << std::defaultfloat << std::setprecision(5)
            << timer() << "s. " << std::endl;
}

template <int Dim>
template <class T>
void PoissonBottomSolver<Dim>::doSolve(Tensor<Real, Dim> &aPhi, const Tensor<Real, Dim> &rhs, const T &bcData) const
{
//  Box<Dim> gbx = emGrid.getGhostedBox();
  // find the homogeneized residual
  GhostFiller<Dim> gfill(emGrid);
//  gfill.fillGhosts(aPhi, bcTypes, bcData);
//  Tensor<Real, Dim> hrsd(gbx);
//  poissonOp.apply(aPhi, rhs, hrsd);
  // flatten the residual
  Tensor<Real, 1> flatRsd(numUnknowns + 1);
  int m = 0;
  ddfor(emGrid, [&](const iVec &idx) {
    if(unknownMap(idx) != -1) {
      flatRsd(m) = rhs(idx); //hrsd(i, j);
      ++m;
    }
  }, false);
  for(; m < numUnknowns + 1; ++m)
    flatRsd(m) = 0.0;
  // solve the linear system
  auto info = LAPACKE_dgetrs(LAPACK_COL_MAJOR, 'N', numActualUnknowns/*n*/, 1/*nrhs*/,
                             poissonFact.data()/*a*/, numUnknowns + 1/*lda*/, ipiv.data()/*ipiv*/,
                             flatRsd.data()/*b*/, numUnknowns + 1/*ldb*/);
  if(info != 0)
    throw std::runtime_error("[PoissonBottomSolver::doSolve] DGETRS fails.");
  // broadcast the solution and re-fill the ghosts
  m = 0;
  ddfor(emGrid, [&](const iVec &idx) {
    if (unknownMap(idx) != -1) {
      aPhi(idx) = flatRsd(m);
      ++m;
    }
  }, false);
  gfill.fillGhosts(aPhi, bcTypes, bcData);
}

template <int Dim>
void PoissonBottomSolver<Dim>::solve(Tensor<Real, Dim> &aPhi, const Tensor<Real, Dim> &rhs, int) const
{
  aPhi = 0.0;
  EmFuncFiller<Dim> eff(emGrid);
  eff.transferBdryValues(aPhi, 0);
  doSolve(aPhi, rhs, 0);
}

//void PoissonBottomSolver<2>::solve(Tensor<Real, 2> &aPhi, const Tensor<Real, 2> &rhs, const Tensor<Real, 1> *bcData) const
//{
//  aPhi = 0.0;
//  EmFuncFiller<Dim> eff(emGrid);
//  eff.transferBdryValues(aPhi, bcData[2*Dim]);
//  doSolve(aPhi, rhs, bcData);
//}

template class PoissonBottomSolver<SpaceDim>;
